import re
import numpy as np
import matplotlib.pyplot as plt

"""
画图函数
对训练过程得到的loss进行画图
"""

# 颜色配置
# colors = ['red','orange','green','cyan','blue','purple','brown','pink','magenta']
colors = ['red', 'orange', 'green', 'blue', 'purple', 'brown', 'pink', 'magenta']


def get_train_info(log_file):
    # 读取文件的训练损失和验证损失值并保存
    train_losses = []
    val_losses = []
    with open(log_file, encoding="utf-8") as f:
        while True:
            lines = f.readline()
            if not lines:
                break
            if len(re.findall(r"(?<=train_loss:)\d+\.?\d*", lines)):
                train_loss = re.findall(r"(?<=train_loss:)\d+\.?\d*", lines)
                train_losses.append(float(train_loss[0]))
            if len(re.findall(r"(?<=val_loss:  )\d+\.?\d*", lines)):
                val_loss = re.findall(r"(?<=val_loss:  )\d+\.?\d*", lines)
                val_losses.append(float(val_loss[0]))
    return train_losses, val_losses


def plat_train_val(title, train_loss, val_loss):
    # 画单个文件的训练损失和验证损失
    plt.figure()
    epochs = len(train_loss)
    x = range(epochs)
    y1 = train_loss
    y2 = val_loss
    plt.plot(x, y1, label='train_loss', color=colors[0])
    plt.plot(x, y2, label='val_loss', color=colors[1])
    plt.title(title)
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend(loc='best')
    plt.show()


def plat_trains(train_dict):
    # 画多个文件的训练损失或者验证损失(保存在字典中)
    plt.figure()
    i = 0
    for key in train_dict.keys():
        x = range(len(train_dict[key]))
        y = train_dict[key]
        plt.plot(x, y, label=key, color=colors[i])
        i += 1
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend(loc='best')
    plt.show()
    # for key in log_dict.keys():
    #     if 'train' in key:
    #         print(log_dict[key])


# if __name__ == '__main__':
#     log_path = "../log/"
#     model_name = "GatedGCN/"
#     test_1 = "test_1.log"
#     test_2 = "test_2.log"
#     train_dict = dict()
#     val_dict = dict()
#     train_dict['test_1'], val_dict['test_1'] = get_train_info(log_path + model_name + test_1)
#     train_dict['test_2'], val_dict['test_2'] = get_train_info(log_path + model_name + test_2)
#     plat_train_val('test_1', train_dict['test_1'], val_dict['test_1'])
#     plat_train_val('test_2', train_dict['test_2'], val_dict['test_2'])
#     plat_trains(train_dict)
#     plat_trains(val_dict)
